# identify_joints.py
import os
import numpy as np
import torch

from NIMBLELayer import NIMBLELayer
import utils
from utils import batch_to_tensor_device

import pytorch3d.io
from pytorch3d.structures.meshes import Meshes

OUTDIR = "joint_debug_output"
os.makedirs(OUTDIR, exist_ok=True)

device = torch.device('cuda')   # 或 'cuda' 若你有 GPU 并已配置好

# --------- load model dicts (跟 demo.py 一样) ----------
pm_dict = np.load("assets/NIMBLE_DICT_9137.pkl", allow_pickle=True)
pm_dict = batch_to_tensor_device(pm_dict, device)

tex_dict = np.load("assets/NIMBLE_TEX_DICT.pkl", allow_pickle=True)
tex_dict = batch_to_tensor_device(tex_dict, device)

nimble_mano_vreg = None
if os.path.exists("assets/NIMBLE_MANO_VREG.pkl"):
    nimble_mano_vreg = np.load("assets/NIMBLE_MANO_VREG.pkl", allow_pickle=True)
    nimble_mano_vreg = batch_to_tensor_device(nimble_mano_vreg, device)

# --------- create layer (关闭 PCA，使用 20 joints × 3 = 60 dims) ----------
nlayer = NIMBLELayer(pm_dict, tex_dict, device,
                     use_pose_pca=False,
                     pose_ncomp=60,   # 这里设置为60更清楚（非 PCA 模式下 forward 期待 20*3）
                     shape_ncomp=20,
                     nimble_mano_vreg=nimble_mano_vreg)

print("=== nlayer summary ===")
print("shape_ncomp:", nlayer.shape_ncomp, "pose_ncomp(attr):", nlayer.pose_ncomp, "tex_ncomp:", nlayer.tex_ncomp)
print("skin_f faces shape:", tuple(nlayer.skin_f.shape))
print("jreg_bone shape:", tuple(nlayer.jreg_bone.shape))
print("jreg_mano shape:", tuple(nlayer.jreg_mano.shape))
print()

# --------- try to print useful constants from utils (if they exist) ----------
print("=== utils module constants (if present) ===")
for name in ["STATIC_JOINT_NUM", "ROOT_JOINT_IDX", "JOINT_PARENT_ID_DICT", "JOINT_ID_BONE_DICT", "JOINT_NAMES"]:
    if hasattr(utils, name):
        print(name, "=", getattr(utils, name))
print()

# --------- baseline (all zeros pose) ----------
bn = 1
pose0 = torch.zeros(bn, 60).to(nlayer.device)   # 20 joints * 3
shape0 = torch.zeros(bn, nlayer.shape_ncomp).to(nlayer.device)
tex0 = torch.zeros(bn, nlayer.tex_ncomp).to(nlayer.device)

# compute baseline bone joints
with torch.no_grad():
    skin0, muscle0, bone0, bone_j0, tex0_img = nlayer.forward(pose0, shape0, tex0, handle_collision=False)
bone_j0 = bone_j0.detach().cpu().numpy()[0]  # shape: (20,3)
print("Baseline bone_joints (20 x 3):\n", bone_j0)
np.savetxt(os.path.join(OUTDIR, "bone_joints_baseline.xyz"), bone_j0)

# --------- perturb each joint (each axis) and measure displacement ----------
angle = 0.5            # 弧度量级，试试 0.3 ~ 1.0 范围
threshold = 1e-5       # 判定为“有影响”的阈值（你可以调大一些）
results = []           # list of tuples (joint_idx, axis, max_displacement, indices_changed)

print("\nPerturbing each joint (20 joints) × 3 axes with angle =", angle)
for j in range(20):
    best_axis = None
    best_max = 0.0
    best_changed_idx = []
    for axis in range(3):
        p = pose0.clone()
        p[0, j*3 + axis] = angle
        with torch.no_grad():
            skin_p, muscle_p, bone_p, bone_j_p, _ = nlayer.forward(p, shape0, tex0, handle_collision=False)
        bone_j_p_np = bone_j_p.detach().cpu().numpy()[0]
        d = np.linalg.norm(bone_j_p_np - bone_j0, axis=1)  # 每个骨关节的位移大小
        maxd = float(d.max())
        changed = np.where(d > threshold)[0].tolist()
        results.append((j, axis, maxd, changed))
        if maxd > best_max:
            best_max = maxd
            best_axis = axis
            best_changed_idx = changed
    print(f"Joint {j:02d}: best axis {best_axis}, max displacement {best_max:.6f}, affected joints: {best_changed_idx}")

# --------- save one mesh per joint (use best axis for each joint) ----------
print("\nSaving .obj for each joint (using each joint's best axis perturbation) ->", OUTDIR)
for j in range(20):
    # find best axis for this joint from results
    rows = [r for r in results if r[0] == j]
    # choose row with max displacement
    if not rows:
        continue
    best = max(rows, key=lambda x: x[2])
    axis = best[1]
    p = pose0.clone()
    p[0, j*3 + axis] = angle
    with torch.no_grad():
        skin_p, _, _, _, _ = nlayer.forward(p, shape0, tex0, handle_collision=False)
    # skin_p: tensor (1, Nskin, 3)
    mesh_p = Meshes(skin_p, nlayer.skin_f.repeat(bn, 1, 1))
    outpath = os.path.join(OUTDIR, f"joint_{j:02d}_axis{axis}.obj")
    pytorch3d.io.IO().save_mesh(mesh_p[0], outpath)
    # also save bone_joints for convenience
    _, _, _, bone_j_p, _ = nlayer.forward(p, shape0, tex0, handle_collision=False)
    np.savetxt(os.path.join(OUTDIR, f"joint_{j:02d}_bonej.xyz"), bone_j_p.detach().cpu().numpy()[0])

print("\nDone. 查看目录 ./", OUTDIR)
print("说明：")
print(" - joint_{i}_axis{a}.obj 可以在 MeshLab/Blender/Windows 3D 查看器里打开，查看哪个手指或掌部发生了位移。")
print(" - joint_{i}_bonej.xyz 列出了 20 个骨关节的 3D 坐标（同 baseline 可以对比）。")
print("\n下一步建议：")
print(" 1) 根据哪个 obj 的手指发生运动，就能把 joint index 标注为 thumb/index/middle/...。")
print(" 2) 得到 index 后，你就可以在脚本里直接对这些 joint 的 axis 分量赋值来实现具体动作（比如拇指伸直其余弯曲）。")
print(" 3) axis 的正负方向和幅度会影响弯曲方向/程度，常用尝试值范围：0.3 ~ 1.0（弧度）。")
